import os
from datetime import datetime
from src.common import commonUtils
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from typing import List, TypedDict, Dict
from langgraph.graph import StateGraph, START, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class Plan(BaseModel):
    index: int = Field(description = "计划序号")
    detail: str = Field(description = "计划详情描述")
    funcName: str = Field(description = "可用的工具名")
    funcParams: Dict[str, str] = Field(description = "json格式的工具的输入参数")

class PlanList(BaseModel):
    plans: List[Plan] = Field(description = "Plan对象的列表")

class ReWOOState(TypedDict):
    query: str
    steps: List[Plan]
    results: dict
    answer: str

class ReWOOAgent:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tools = agent_tools

    def plan_node(self, state: ReWOOState):
        query = state["query"]
        template = """
            当前时间是：{current_date}
        
            针对以下任务，制定分步骤解决问题的计划。不要添加任何多余的步骤。
            针对每个计划，需明确指定用于获取证据的外部工具及其输入参数。
            可将证据存储到变量#E中，以便后续工具调用。
            （计划，#E1，计划，#E2，计划，...）
                
            工具可以是以下之一：{tool_descs}
                
            例如，
            任务：托马斯、托比和丽贝卡一周共工作了157小时。
            Thomas工作了x小时。
            托比的工作时间比托马斯少10小时，是托马斯工作时间的两倍，而丽贝卡比托比少8小时。
            丽贝卡工作了多少小时？
                
            计划：假设托马斯工作了x小时，将问题转化为代数表达式，并使用Wolfram Alpha求解。#E1=WolframAlpha[Solve x+（2x−10）+（（2x−10）−8）=157]
            计划：计算托马斯的工作时长。#E2=LLM[根据#E1求解x]
            计划：计算丽贝卡工作的时间。#E3=Calculator[（2∗#E2−10）−8]
                
            开始！请详细描述您的计划。每个计划只能跟随一个#E，并严格按格式输出计划。
                
            任务: {task}
        """
        tool_descs = commonUtils.build_tools_description(self.agent_tools)
        current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        prompt = template.format(current_date = current_date, tool_descs = tool_descs, task = query)
        prompt = ChatPromptTemplate.from_messages([("user", prompt)])
        structure_llm = self.llm_model.with_structured_output(PlanList)
        response = (prompt | structure_llm).invoke({ "task": query })
        logger.info(f"ReWOOAgent plan_node response: {response}")
        return { "steps": response.plans }

    def get_current_task(self, state: ReWOOState):
        results = state.get("results", None)
        if results is None:
            return 1

        if len(results) == len(state["steps"]):
            return None
        else:
            return len(results) + 1

    def tool_node(self, state):
        step = self.get_current_task(state)
        logger.info(f"ReWOOAgent tool_node step: {step}")
        steps = state.get("steps", None)
        plan = steps[step - 1]
        logger.info(f"ReWOOAgent tool_node plan: {plan}")

        results = state.get("results", {})
        for tool in self.agent_tools:
            if plan.funcName == tool.name:
                result = tool.func(**plan.funcParams)
                results[plan.index] = result
        return { "results": results }

    def solve_node(self, state):
        plan = ""
        steps = state["steps"]
        for index, detail, funcName, funcParams in steps:
            plan += f"Plan: {index} {detail} = {funcName}[{funcParams}]"

        template = """
            当前时间是：{current_date}
        
            解决以下任务或问题。
            为解决该问题已分步骤制定计划，并为每个计划获取了对应的证据。
            使用时需谨慎，因为长篇证据可能包含无关信息。
                
            {plan}
                
            现在根据上面提供的证据来解决这个问题或任务。
            请直接回答，无需额外说明。
                
            任务: {task}
                
            回复：
        """
        prompt = template.format(current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S"), plan = plan, task = state["query"])
        response = (self.llm_model | StrOutputParser()).invoke(prompt)
        return { "answer": response }

    def route_node(self, state):
        step = self.get_current_task(state)
        if step is None:
            return "solve"
        else:
            return "tool"

    def build_graph(self):
        graph = StateGraph(ReWOOState)
        graph.add_node("plan", self.plan_node)
        graph.add_node("tool", self.tool_node)
        graph.add_node("solve", self.solve_node)

        graph.add_edge(START, "plan")
        graph.add_edge("plan", "tool")
        graph.add_conditional_edges("tool", self.route_node, { "solve": "solve", "tool": "tool" })
        graph.add_edge("solve", END)

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agent_ReWOO_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"ReWOOAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        answer = response.get("answer", None)
        logger.info(f"ReWOOAgent invoke answer len: {len(str(answer))}")
        return { "chain_result": answer }
